package com.sh.data.engine.domain.shims.influxdb.manager;

import com.google.common.collect.Lists;
import com.influxdb.client.InfluxDBClient;
import com.influxdb.client.InfluxDBClientFactory;
import com.influxdb.client.InfluxDBClientOptions;
import com.influxdb.client.InfluxQLQueryApi;
import com.influxdb.client.domain.Bucket;
import com.influxdb.client.domain.InfluxQLQuery;
import com.influxdb.query.InfluxQLQueryResult;
import com.sh.data.engine.domain.shims.db.model.FieldInfoDomain;
import com.sh.data.engine.domain.shims.db.model.PreviewDataDomain;
import com.sh.data.engine.domain.shims.db.model.TableInfoDomain;
import com.sh.data.engine.domain.shims.influxdb.option.InfluxDbOptions;
import lombok.extern.slf4j.Slf4j;
import okhttp3.OkHttpClient;
import org.springframework.util.CollectionUtils;

import java.sql.SQLException;
import java.util.*;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

/**
 * @author: mengzheng.mhc
 * @date: 2024/6/11 18:13
 */
@Slf4j
public class InfluxDBManager {

    private static final String TIME_FIELD_TYPE = "time";

    private static final String TAG_VALUE_TYPE = "string";

    private static final Integer DEFAULT_PREVIEW_SIZE = 10;

    private final InfluxDbOptions options;

    public InfluxDBManager(InfluxDbOptions options) {
        this.options = options;
    }

    public void testConnection() throws SQLException {
        InfluxDBClientOptions influxDBClientOptions = InfluxDBClientOptions.builder()
                .url(options.getUrl())
                .authenticateToken(options.getApiToken().toCharArray())
                .org(options.getOrgName())
                .okHttpClient(new OkHttpClient.Builder().readTimeout(1, TimeUnit.MINUTES))
                .build();
        InfluxDBClient influxDBClient =
            InfluxDBClientFactory.create(influxDBClientOptions);

        // 首先判定能否ping通，确保url正确
        Boolean result = influxDBClient.ping();
        if (Boolean.TRUE.equals(result)) {
            // PING请求成功，验证认证信息和Bucket是否有效
            try {
                result = checkDatabaseExists(influxDBClient, options.getDbName());
                if (Boolean.TRUE.equals(result)) {
                    return;
                }
            } catch (Exception ex) {
                log.error("InfluxDB ping request wasn't successful", ex);
            } finally {
                influxDBClient.close();
            }
        }

        throw new SQLException("InfluxDB ping request wasn't successful");
    }

    public List<Map<String, Object>> list(String query) {
        List<InfluxQLQueryResult.Series> seriesList = querySeries(query);
        if (CollectionUtils.isEmpty(seriesList)) {
            return Collections.emptyList();
        }

        return seriesList.stream()
            .flatMap(series -> {
                List<String> fieldNameList = series.getColumns().entrySet().stream()
                    .sorted(Map.Entry.comparingByValue())
                    .map(Map.Entry::getKey)
                    .collect(Collectors.toList());
                return series.getValues().stream()
                    .map(value -> {
                        Map<String, Object> row = new LinkedHashMap<>();
                        for (int i = 0; i < fieldNameList.size(); i++) {
                            row.put(fieldNameList.get(i), value.getValues()[i]);
                        }
                        series.getTags().forEach(row::putIfAbsent);
                        return row;
                    });
            })
            .collect(Collectors.toList());
    }

    public List<Map<String, Object>> list(String query, Integer pageNo, Integer pageSize) {
        String paginatedQuery =
            "SELECT * FROM (" + query + ") LIMIT " + pageSize + " OFFSET " + (pageNo - 1) * pageSize;
        return list(paginatedQuery);
    }

    public Long count(String query) {
        String countQuery = "SELECT COUNT(*) FROM (" + query + ")";
        InfluxQLQueryResult.Series series = querySingleSeries(countQuery);
        if (series == null) {
            return 0L;
        }

        // COUNT查询的返回结果中会包含time列和count_[fieldKey]列，所以要获取count_[fieldKey]列的实际名称
        String countColumn =
            series.getColumns().keySet().stream()
                .filter(s -> !s.equalsIgnoreCase("time"))
                .findFirst()
                .get();
        String countStringValue = series.getValues().get(0).getValueByKey(countColumn).toString();
        return Long.parseLong(countStringValue);
    }

    public List<TableInfoDomain> getTableAndViewList(String dbName) {
        List<TableInfoDomain> tables = new ArrayList<>();
        InfluxQLQueryResult.Series series = querySingleSeries("SHOW MEASUREMENTS ON " + escapeIdentifier(dbName));
        if (series == null) {
            return tables;
        }

        String columnName = series.getColumns().keySet().iterator().next();
        for (InfluxQLQueryResult.Series.Record value : series.getValues()) {
            String measurement = value.getValueByKey(columnName).toString();
            TableInfoDomain tableInfoDomain = new TableInfoDomain();
            tableInfoDomain.setTableName(measurement);
            tableInfoDomain.setDbName(dbName);
            tableInfoDomain.setTblType(TableInfoDomain.TblType.MANAGED_TABLE);
            tables.add(tableInfoDomain);
        }

        return tables;
    }

    public PreviewDataDomain previewData(String tableName) {
        return previewData(tableName, DEFAULT_PREVIEW_SIZE);
    }

    public PreviewDataDomain previewData(String tableName, Integer pageSize) {
        InfluxQLQueryResult.Series series =
            querySingleSeries("SELECT * FROM " + escapeIdentifier(tableName) + " LIMIT " + pageSize);
        if (series == null) {
            return null;
        }

        List<String> fieldNameList =
            series.getColumns().entrySet().stream()
                .sorted(Map.Entry.comparingByValue())
                .map(Map.Entry::getKey)
                .collect(Collectors.toList());

        List<List<Object>> dataList = new ArrayList<>();
        for (InfluxQLQueryResult.Series.Record value : series.getValues()) {
            List<Object> data = Arrays.asList(value.getValues());
            dataList.add(data);
        }

        return PreviewDataDomain.builder().fieldNameList(fieldNameList).dataList(dataList).build();
    }

    public List<FieldInfoDomain> getFieldList(String dbName, String tableName) {
        // 之前按照time、tag keys、field key的顺序返回所有field，这个逻辑在某条记录是多field的时候不成立
        // 所以现在改成按照previewData的逻辑，尝试获取一条数据，然后依据这条记录返回结果里面的顺序来返回

        // time列表
        FieldInfoDomain timeFieldInfo =
            FieldInfoDomain.builder().fieldName("time").fieldType(TIME_FIELD_TYPE).isTag(false).build();

        // tag列
        Map<String, FieldInfoDomain> tagColumns = new HashMap<>();
        InfluxQLQueryResult.Series series = querySingleSeries("SHOW TAG KEYS FROM " + escapeIdentifier(tableName));
        if (series != null) {
            // 可能不存在tag列
            for (InfluxQLQueryResult.Series.Record value : series.getValues()) {
                String influxTagKey = value.getValueByKey("tagKey").toString();
                FieldInfoDomain influxTagInfo =
                    FieldInfoDomain.builder()
                        .fieldName(influxTagKey)
                        .fieldType(TAG_VALUE_TYPE)
                        .isTag(true)
                        .build();
                tagColumns.put(influxTagKey, influxTagInfo);
            }
        }

        series = querySingleSeries("SHOW FIELD KEYS FROM " + escapeIdentifier(tableName));
        if (series == null) {
            return null;
        }

        // field列
        Map<String, FieldInfoDomain> fieldColumns = new HashMap<>();
        for (InfluxQLQueryResult.Series.Record value : series.getValues()) {
            String influxFieldKey = value.getValueByKey("fieldKey").toString();
            String influxFieldType = value.getValueByKey("fieldType").toString();
            FieldInfoDomain influxFieldInfo =
                FieldInfoDomain.builder()
                    .fieldName(influxFieldKey)
                    .fieldType(influxFieldType)
                    .isTag(false)
                    .build();
            fieldColumns.put(influxFieldKey, influxFieldInfo);
        }

        PreviewDataDomain previewDataDomain = previewData(tableName, 1);
        if (previewDataDomain == null) {
            return null;
        }

        List<FieldInfoDomain> fieldInfos = Lists.newArrayList();
        List<String> fieldNameList = previewDataDomain.getFieldNameList();
        for (String fieldName : fieldNameList) {
            if (fieldName.equals("time")) {
                fieldInfos.add(timeFieldInfo);
            } else if (tagColumns.containsKey(fieldName)) {
                fieldInfos.add(tagColumns.get(fieldName));
            } else if (fieldColumns.containsKey(fieldName)) {
                fieldInfos.add(fieldColumns.get(fieldName));
            }
        }

        return fieldInfos;
    }

    private String escapeIdentifier(String identifier) {
        // Check if the identifier is already escaped
        if (identifier.startsWith("\"") && identifier.endsWith("\"")) {
            return identifier;
        }

        return "\"" + identifier + "\"";
    }

    private List<InfluxQLQueryResult.Series> querySeries(String command) {
        InfluxDBClient influxDBClient = null;

        try {
            influxDBClient =
                InfluxDBClientFactory.create(
                    options.getUrl(), options.getApiToken().toCharArray(), options.getOrgName());

            InfluxQLQueryApi queryApi = influxDBClient.getInfluxQLQueryApi();
            InfluxQLQuery query = new InfluxQLQuery(command, options.getDbName());
            InfluxQLQueryResult queryResult = queryApi.query(query);
            if (queryResult == null) {
                return null;
            }

            List<InfluxQLQueryResult.Result> results = queryResult.getResults();
            if (results == null) {
                return null;
            }

            InfluxQLQueryResult.Result result = results.get(0);
            if (result == null) {
                return null;
            }

            return result.getSeries();
        } finally {
            if (influxDBClient != null) {
                influxDBClient.close();
            }
        }
    }

    private InfluxQLQueryResult.Series querySingleSeries(String command) {
        List<InfluxQLQueryResult.Series> series = querySeries(command);
        return series != null && !series.isEmpty() ? series.get(0) : null;
    }

    private boolean checkDatabaseExists(InfluxDBClient influxDBClient, String dbName) {
        Bucket bucket = influxDBClient.getBucketsApi().findBucketByName(dbName);
        return bucket != null;
    }

    public static void main(String[] args) {
        InfluxDbOptions influxDbOptions =
            InfluxDbOptions.builder()
                .dbName("sisfirststage")
                .apiToken(
                    "sS_6Bt_k0_HC5CZME7n1ukVYGhoyTFI8CU2ODBRHwT1CGKCB6xjG72ccZcAjTn1AG5C5L0Thi3bdYYgl2JpF9g==")
                .orgName("xuelang")
                .url("http://osdev.xuelangyun.com:30198")
                .build();
        InfluxDBManager influxDBManager = new InfluxDBManager(influxDbOptions);
        try {
            influxDBManager.testConnection();
        } catch (SQLException e) {
            System.out.println("test connection failed = " + e);
            return;
        }

        List<TableInfoDomain> tableAndViewList = influxDBManager.getTableAndViewList("sisfirststage");
        System.out.println("tableAndViewList = " + tableAndViewList);

        PreviewDataDomain previewDataDomain = influxDBManager.previewData("mz2");
        System.out.println("previewDataDomain = " + previewDataDomain.getFieldNameList());

        List<FieldInfoDomain> fieldList = influxDBManager.getFieldList("sisfirststage", "mz2");
        System.out.println(
            "fieldList = "
                + fieldList.stream().map(FieldInfoDomain::getFieldName).collect(Collectors.toList()));

        List<Map<String, Object>> multipleList =
            // select * from sis4 where pointId = '${pid}'
            influxDBManager.list("select count(cpu) from mz2 group by region");
        System.out.println("multipleList = " + multipleList);

        List<Map<String, Object>> list =
            // select * from sis4 where pointId = '${pid}'
            influxDBManager.list("select * from sis4 where pointId = '10'");
        System.out.println("list = " + list);

        Long count = influxDBManager.count("SELECT * from sis4 where pointId = '10'");
        System.out.println("count = " + count);

        List<Map<String, Object>> paginatedList =
            influxDBManager.list("select * from sis4 where pointId = '10'", 1, 3);
        System.out.println("paginatedList = " + paginatedList);
    }
}
